package com.tony.ai.camerax

import android.annotation.SuppressLint
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.ImageFormat
import android.graphics.YuvImage
import android.media.Image
import androidx.camera.core.ImageProxy
import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer

object ImageUtils {

    /**
     * 将ImageProxy（YUV_420_888格式）转换为Bitmap
     *
     * @param imageProxy 输入的ImageProxy对象
     * @return 转换后的Bitmap对象
     */
    @SuppressLint("UnsafeExperimentalUsageError")
    fun imageProxyToBitmap(imageProxy: ImageProxy): Bitmap? {
        if (imageProxy.format != ImageFormat.YUV_420_888) {
            throw IllegalArgumentException("Invalid image format, expected YUV_420_888")
        }
        val image = imageProxy.image!!

        val yBuffer = image.planes[0].buffer
        val uBuffer = image.planes[1].buffer
        val vBuffer = image.planes[2].buffer

        val ySize = yBuffer.remaining()
        val uSize = uBuffer.remaining()
        val vSize = vBuffer.remaining()

        val nv21 = ByteArray(ySize + uSize + vSize)

        // Y buffer is first
        yBuffer.get(nv21, 0, ySize)

        // V buffer is second
        vBuffer.get(nv21, ySize, vSize)

        // U buffer is third
        uBuffer.get(nv21, ySize + vSize, uSize)

        val yuvImage = YuvImage(nv21, ImageFormat.NV21, image.width, image.height, null)
        val out = ByteArrayOutputStream()
        yuvImage.compressToJpeg(android.graphics.Rect(0, 0, image.width, image.height), 100, out)

        val imageBytes = out.toByteArray()
        return BitmapFactory.decodeByteArray(imageBytes, 0, imageBytes.size)
    }

    fun yuv420888ToBitmap(image: Image): Bitmap {
        val width = image.width
        val height = image.height

        // 获取 YUV 数据的三个平面
        val planes: Array<Image.Plane> = image.planes
        val yBuffer: ByteBuffer = planes[0].buffer
        val uBuffer: ByteBuffer = planes[1].buffer
        val vBuffer: ByteBuffer = planes[2].buffer

        // Y 数据的步长
        val yPixelStride = planes[0].pixelStride
        val yRowStride = planes[0].rowStride
        val uvPixelStride = planes[1].pixelStride
        val uvRowStride = planes[1].rowStride

        // 创建 ARGB_8888 格式的 Bitmap
        val bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
        val pixels = IntArray(width * height)

        // YUV 到 RGB 的转换
        for (y in 0 until height) {
            for (x in 0 until width) {
                // Y 分量
                val yValue = yBuffer.get(y * yRowStride + x * yPixelStride).toInt() and 0xff

                // U 和 V 分量
                val uvX = x / 2
                val uvY = y / 2
                val uValue = uBuffer.get(uvY * uvRowStride + uvX * uvPixelStride).toInt() and 0xff
                val vValue = vBuffer.get(uvY * uvRowStride + uvX * uvPixelStride).toInt() and 0xff

                // YUV 到 RGB 转换公式
                val r = (yValue + 1.402 * (vValue - 128)).toInt()
                val g = (yValue - 0.34414 * (uValue - 128) - 0.71414 * (vValue - 128)).toInt()
                val b = (yValue + 1.772 * (uValue - 128)).toInt()

                // 限制 RGB 值在 0-255 之间
                val clampedR = r.coerceIn(0, 255)
                val clampedG = g.coerceIn(0, 255)
                val clampedB = b.coerceIn(0, 255)

                // 将 RGB 转换为 ARGB
                pixels[y * width + x] = 0xff000000.toInt() or (clampedR shl 16) or (clampedG shl 8) or clampedB
            }
        }

        // 将像素数据写入 Bitmap
        bitmap.setPixels(pixels, 0, width, 0, 0, width, height)

        return bitmap
    }
}